from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, Add, Concatenate, Layer
from tensorflow.keras.optimizers import Adam
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv
import tempfile
import zipfile
import os
# Add noise to images (Salt & Pepper noise for demonstration)
def add_noise(img, noise_type="s&p"):
noise = np.random.rand(*img.shape)
noisy_img = img.copy()
noisy_img[noise < 0.1] = 0 # Add black pixels
noisy_img[noise > 0.9] = 1 # Add white pixels
return noisy_img
# Apply Gaussian blur with much stronger blur for the first image (extra blur)
def apply_blur(img, num_stages=7):
blurred_images = []
# Apply an extremely strong blur for the first image (making it very misty and blurry)
first_image_blur_strength = 2 * num_stages + 15 # Very large kernel size for first image
first_blurred_img = cv.GaussianBlur(img, (first_image_blur_strength, first_image_blur_strength), 0)
# Apply an additional blur pass to make it even more blurry (extreme blur effect)
first_blurred_img = cv.GaussianBlur(first_blurred_img, (first_image_blur_strength, first_image_blur_strength), 0)
blurred_images.append(first_blurred_img)
# Apply decreasing blur for the rest of the images
for i in range(1, num_stages):
blur_strength = 2 * (num_stages - i) + 1 # Gradually reduce the blur
blurred_img = cv.GaussianBlur(img, (blur_strength, blur_strength), 0)
blurred_images.append(blurred_img)
return blurred_images
# Preprocess images taken from multiple angles
def preprocess_multi_angle_images(zip_file):
angles = []
with tempfile.TemporaryDirectory() as temp_dir:
with zipfile.ZipFile(zip_file, 'r') as zip_ref:
zip_ref.extractall(temp_dir)
for file in sorted(os.listdir(temp_dir)): # Sorting ensures consistent processing
img_path = os.path.join(temp_dir, file)
if file.lower().endswith(('.png', '.jpg', '.jpeg')):
img = cv.imread(img_path, cv.IMREAD_UNCHANGED)
img = cv.cvtColor(img, cv.COLOR_BGR2RGB) / 255.0
img_resized = cv.resize(img, (256, 256))
angles.append(img_resized)
return np.array(angles) # Shape: (num_angles, height, width, channels)
# Custom Layer for Feature Aggregation
class FeatureAggregationLayer(Layer):
def __init__(self, num_angles, **kwargs):
super(FeatureAggregationLayer, self).__init__(**kwargs)
self.num_angles = num_angles
def build(self, input_shape):
# Create a learnable weight matrix for feature aggregation
self.conv = Conv2D(64, (1, 1), activation='relu', padding='same', use_bias=False)
super(FeatureAggregationLayer, self).build(input_shape)
def call(self, inputs):
# Concatenate features from different angles
aggregated = Concatenate(axis=-1)(inputs)
return self.conv(aggregated)
# Bracketing image restoration and enhancement model for multi-angle inputs
def build_multi_angle_model(input_shape=(256, 256, 3), num_angles=7):
inputs = [Input(shape=input_shape) for _ in range(num_angles)]
angle_features = []
for inp in inputs:
conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(inp)
conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv1)
angle_features.append(conv2)
# Use the custom feature aggregation layer
aggregated = FeatureAggregationLayer(num_angles)(angle_features)
# Enhancement and restoration with residual connections
conv3 = Conv2D(64, (3, 3), activation='relu', padding='same')(aggregated)
conv4 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv3)
restored_output = Conv2D(3, (3, 3), activation='relu', padding='same')(conv4) # Changed to ReLU activation
model = Model(inputs=inputs, outputs=restored_output)
model.compile(optimizer=Adam(learning_rate=0.001), loss='mse')
return model
# Process multi-angle images and predict
def predict_multi_angle_images(zip_file, model, num_angles=7):
angles = preprocess_multi_angle_images(zip_file)
if len(angles) < num_angles:
raise ValueError(f"The model expects {num_angles} angles, but fewer were provided.")
if len(angles) > num_angles:
angles = angles[:num_angles] # Use only the first `num_angles` inputs
noisy_inputs = [add_noise(angle) for angle in angles] # Add noise to inputs
# Expand dimensions for each input to match the model's expected shape
inputs = [np.expand_dims(noisy, axis=0) for noisy in noisy_inputs]
# Generate blurred versions of each image in the dataset (starting with strong blur for the first image)
all_blurred_images = [] # To store blurred images for each of the 7 images
for angle in angles:
blurred_images = apply_blur(angle, num_stages=num_angles) # Apply to each image in the dataset
all_blurred_images.append(blurred_images)
# Model prediction (restoration of the image)
restored_image = model.predict(inputs)
# Ensure proper scaling for the final restored image
restored_image_rescaled = np.clip(restored_image[0] * 255, 0, 255).astype(np.uint8)
# Plot results for each of the 7 images
plt.figure(figsize=(20, 12))
# Show the blurred transitions for each image (starting with heavy blur, gradually less)
for img_idx, blurred_images in enumerate(all_blurred_images):
for i, blurred_img in enumerate(blurred_images):
plt.subplot(len(angles), num_angles, img_idx * num_angles + i + 1)
plt.imshow(blurred_img)
if img_idx == 0:
plt.title(f"Image {img_idx+1} - Blurred {i+1}")
plt.axis('off')
plt.show()
# Build and use the model
multi_angle_model = build_multi_angle_model()
zip_file_path = 'test_A.zip' # Replace with the actual zip file path
predict_multi_angle_images(zip_file_path, multi_angle_model)
